﻿#include "mediancutquantizer.h"

namespace WingGif {

MedianCutQuantizer::MedianCutCube::MedianCutCube(QList<QColor> colors) {
  _colorList = colors;
  shrink();
}

int MedianCutQuantizer::MedianCutCube::redSize() {
  return _redHighBound - _redLowBound;
}

int MedianCutQuantizer::MedianCutCube::greenSize() {
  return _greenHighBound - _greenLowBound;
}

int MedianCutQuantizer::MedianCutCube::blueSize() {
  return _blueHighBound - _blueLowBound;
}

QColor MedianCutQuantizer::MedianCutCube::color() {
  if (_cubeColor.isValid())
    return _cubeColor;

  int red = 0, green = 0, blue = 0;

  for (auto color : _colorList) {
    red += color.red();
    green += color.green();
    blue += color.blue();
  }

  auto colorsCount = _colorList.count();

  if (colorsCount != 0) {
    red /= colorsCount;
    green /= colorsCount;
    blue /= colorsCount;
  }
  _cubeColor = QColor::fromRgb(red, green, blue);
  return _cubeColor;
}

void MedianCutQuantizer::MedianCutCube::shrink() {
  // Get the minimum/maximum values for each RGB component of specified colors.
  _redLowBound = _greenLowBound = _blueLowBound = 255;
  _redHighBound = _greenHighBound = _blueHighBound = 0;

  for (auto colort : _colorList) {
    if (colort.red() < _redLowBound)
      _redLowBound = uchar(colort.red());
    if (colort.red() > _redHighBound)
      _redHighBound = uchar(colort.red());

    if (colort.green() < _greenLowBound)
      _greenLowBound = uchar(colort.green());
    if (colort.green() > _greenHighBound)
      _greenHighBound = uchar(colort.green());

    if (colort.blue() < _blueLowBound)
      _blueLowBound = uchar(colort.blue());
    if (colort.blue() > _blueHighBound)
      _blueHighBound = uchar(colort.blue());
  }
}

void MedianCutQuantizer::MedianCutCube::splitAtMedian(
    uchar componentIndex, MedianCutQuantizer::MedianCutCube *&medianCube1,
    MedianCutQuantizer::MedianCutCube *&medianCube2) {
  // 原来的代码这三种情况都是第一个代码内容，感觉十分奇怪，我给改了
  switch (componentIndex) {
  case 0: {
    std::sort(_colorList.begin(), _colorList.end(),
              [=](QColor &p, QColor &n) { return p.red() < n.red(); });
    break;
  }
  case 1: {
    std::sort(_colorList.begin(), _colorList.end(),
              [=](QColor &p, QColor &n) { return p.green() < n.green(); });
    break;
  }
  case 2: {
    std::sort(_colorList.begin(), _colorList.end(),
              [=](QColor &p, QColor &n) { return p.blue() < n.blue(); });
    break;
  }
  }

  auto medianIndex = _colorList.count() >> 1;

  medianCube1 = new MedianCutCube(_colorList.mid(0, medianIndex));
  medianCube2 = new MedianCutCube(
      _colorList.mid(medianIndex, _colorList.count() - medianIndex));
}

void MedianCutQuantizer::MedianCutCube::setPaletteIndex(int newPaletteIndex) {
  paletteIndex = newPaletteIndex;
}

bool MedianCutQuantizer::MedianCutCube::isColorIn(QColor color) {
  return (color.red() >= _redLowBound && color.red() <= _redHighBound) &&
         (color.green() >= _greenLowBound &&
          color.green() <= _greenHighBound) &&
         (color.blue() >= _blueLowBound && color.blue() <= _blueHighBound);
}

MedianCutQuantizer::MedianCutQuantizer() : Quantizer(false, this) {}

void MedianCutQuantizer::initialQuantizePixel(QColor pixel) {
  if (pixel.alpha() == 0)
    return;

  colors.append(pixel);
}

QVector<QColor> MedianCutQuantizer::buildPalette() {
  auto _transparentColor = transparentColor();
  auto _maxColors = maxColors();
  auto _maxColorsWithTransparency =
      _transparentColor.isValid() ? _maxColors - 1 : _maxColors;
  setMaxColorsWithTransparency(_maxColorsWithTransparency);

  // Quantization.
  _cubes.append(new MedianCutCube(colors));

  // Split the cube until we get required amount of colors.
  splitCubes(_cubes, _maxColorsWithTransparency);

  // Get the final palette.
  auto palette = QVector<QColor>(_maxColors);

  for (auto i = 0; i < _maxColorsWithTransparency; i++) {
    palette.append(_cubes[i]->color());
    _cubes[i]->setPaletteIndex(i);
  }

  // Add the transparent color to the last position.
  if (_transparentColor.isValid())
    palette.append(QColor::fromRgba(qRgba(_transparentColor.red(),
                                          _transparentColor.green(),
                                          _transparentColor.blue(), 0)));

  return palette;
}

uchar MedianCutQuantizer::quantizePixel(QColor pixel) {
  QList<MedianCutCube *> res;
  for (auto item : _cubes) {
    if (item->isColorIn(pixel))
      res.append(item);
  }

  for (auto &cube : res)
    return uchar(cube->paletteIndex);

  return 0;
}

void MedianCutQuantizer::splitCubes(
    QList<MedianCutQuantizer::MedianCutCube *> &cubes, int count) {
  auto cubeIndexToSplit = cubes.count() - 1;

  while (cubes.count() < count) {
    auto cubeToSplit = cubes[cubeIndexToSplit];
    MedianCutCube *cube1, *cube2;

    // Find the longest color size to use for splitting.
    if (cubeToSplit->redSize() >= cubeToSplit->greenSize() &&
        cubeToSplit->redSize() >= cubeToSplit->blueSize())
      cubeToSplit->splitAtMedian(0, cube1, cube2);
    else if (cubeToSplit->greenSize() >= cubeToSplit->blueSize())
      cubeToSplit->splitAtMedian(1, cube1, cube2);
    else
      cubeToSplit->splitAtMedian(2, cube1, cube2);

    // Remove the old "big" cube.
    cubes.removeAt(cubeIndexToSplit);

    // Add two smaller cubes instead
    cubes.insert(cubeIndexToSplit, cube1);
    cubes.insert(cubeIndexToSplit, cube2);

    if (--cubeIndexToSplit < 0)
      cubeIndexToSplit = cubes.count() - 1;
  }
}

} // namespace WingGif
